# -*- coding:utf-8 -*-
'''
    Implement HybridNode class for the latest proposed code struct.
    HybridNode can reconstruct internal nodes and the connection around themselves.
'''

# May it be a cnode class?

from  builtins import range
from functools import reduce
import numpy as np
from cnode import *
from cmsg import *
import logging
import cgraph
from satProp import *
from motionModel import satellitePropMotion,staticMotion

class cHybridNode(cNode):
    '''
        hybrid node class, son class of cNode
        graph: inner graph, it is a cGraph entity
        defaultlinkupTime: default linkup period, relative to current time
        hasMutualMeasure : whether there is a mutual measure
        backwardMode: enable backStage 3
        TODO currently, we don't support emit/receive multiple links simutanously.
    '''
    def __init__(self,name=None,backwardMode=None):
        if name is None:
            name='ahybridNode'
        cNode.__init__(self,name)
        self.type = 'hybrid'
        self.graph= None
        self.hasMutualMeasure= False
        self.passMutualMeasure=False # whether should pass a mutual measure.
        if backwardMode is None:
            backwardMode =1 
        self.backwardMode = backwardMode


    def constructGraph(self,para):
        '''
            construct inner graph with specific parameters
            para: dict-type parameters setting
                'hasMutualMeasure': boolean, indicate whether there is a mutual measurement.
                'initialStatement': initial measurement is key parameters to trans, including:(is hasMutualMeasure)
                    'emitKineMsg': kineMsg of emit sat
                    'recKineMsg':  kineMsg of receive satMsg
                    'emitPosMsg': Position Msg of emit sat
                    'recPosMsg' : Position Msg of rec sat
                    'emitTimeDriftMsg': DriftMsg 
                    'recTimeDriftMsg': DriftMsg 
                    'link1startTime': 
                    'link1endTime'  :  
                    'link1measurement': the measurement for link1 
                    'link2startTime':
                    'link2endTime':
                    'link2measurement', similiar to above
                    'emitSat': 
                    'recSat'
        '''
        self.graph = cgraph.cGraph()
        self.hasMutualMeasure = para['hasMutualMeasure']
        # constuct graph structure
        nodes     = [] 
        nodes.append( priorFactNode('PriorMobilej1'))
        nodes.append( agentVarNode('Posteriorj1'))
        if(self.hasMutualMeasure):
            nodes.append(agentVarNode('Posteriori1'))
            nodes.append(agentVarNode('Posteriori2'))
            nodes.append(agentVarNode('Posteriorj2'))
            nodes.append(priorFactNode('PriorMobilei1'))

            nodes.append(mobilityFactNode('PriorMobilei2'))
            nodes.append(mobilityFactNode('PriorMobilej2'))
            nodes.append(priorFactNode('TimeDiffPrior'))
            nodes.append(mutualConWithDriftNode('MutualProp1'))
            nodes.append(mutualConWithDriftNode('MutualProp2'))
            nodes.append(timediffVarNode('TimeDiff'))
            self.graph.addNodes(nodes)

            self.graph.connectVarNodes('Posteriori1','PriorMobilei1','in')
            self.graph.connectVarNodes('Posteriorj1','PriorMobilej1','in')

            self.graph.connectVarNodes(['Posteriori1','Posteriori2'],'PriorMobilei2',['out','in'])
            self.graph.connectVarNodes(['Posteriorj1','Posteriorj2'],'PriorMobilej2',['out','in'])

            self.graph.connectVarNodes(['Posteriori1','Posteriorj1'],'MutualProp1',['out','in'])
            self.graph.connectVarNodes(['Posteriorj2','Posteriori2'],'MutualProp2',['out','in'])
            self.graph.connectVarNodes('TimeDiff','MutualProp1',['out'])
            self.graph.connectVarNodes('TimeDiff','MutualProp2',['out'])
            self.graph.connectVarNodes('TimeDiff','TimeDiffPrior',['in'])

        else:
            self.graph.addNodes(nodes)
            self.graph.connectVarNodes('Posteriorj1','PriorMobilej1','in')

        # construct Msg therein, because some msg should be assigned in there
        self.constructPriori(para['initialStatement'])
    
    def constructPriori(self,para):
        if self.hasMutualMeasure:
            # priorMsg
            gins = self.graph.getNodeInstance 
            gins('PriorMobilei1').setPriorMsg(para['emitKineMsg'])
            gins('PriorMobilej1').setPriorMsg(para['recKineMsg'])
            timeDriftMsg  =  para['emitTimeDriftMsg'] - para['emitTimeDriftMsg']   
            gins('TimeDiffPrior').setPriorMsg(timeDriftMsg)
            # motionmodel

            gins('PriorMobilei2').setMobilityModel(satellitePropMotion())
            gins('PriorMobilej2').setMobilityModel(satellitePropMotion())
            gins('PriorMobilei2').model.updatePara({'startTime':para['link1startTime'],
                                              'endTime':para['link2endTime'],'satname':para['emitSatName']})
            gins('PriorMobilej2').model.updatePara({'startTime':para['link1endTime'],
            'endTime':para['link2startTime'],'satname':para['recSatName']})
            
            gins('PriorMobilei2').model.setForward()
            gins('PriorMobilej2').model.setForward()

            
            #mutualConNode
            gins('MutualProp1').setMeasuremnt(para['link1measurement'],DEFAULTMEARSUREMENTVAR)
            gins('MutualProp1').setFirstLink()
            gins('MutualProp2').setMeasuremnt(para['link2measurement'],DEFAULTMEARSUREMENTVAR)
            gins('MutualProp2').setFirstLink(False)

    def getMutualMeasureTime(self):
        # get mutual measurement time of my clock
        pass

    def run(self,loopTime=1):
        '''
            run the node to generate the node posterior and pass msgs
            loopTime: how many times should do loop
        '''

        gins = self.graph.getNodeInstance
        
        # sat1covTrs  =  np.zeros((loopTime+1))
        # sat2covTrs  =  np.zeros((loopTime+1))
        # timediff  =   np.zeros((loopTime+1))
        self.forwardPass()
        # sat1covTrs[0]= np.trace(gins('Posteriori2').processData(0).cov[0:2,0:2])
        # sat2covTrs[0]= np.trace(gins('Posteriorj2').processData(0).cov[0:2,0:2])
        # timediff[0]  = gins('TimeDiff').processData(0).cov
        for loopIndex in range(0,loopTime):
            self.backwardPass(self.backwardMode)
            self.forwardPass()
            # sat1covTrs[loopIndex+1]= np.trace(gins('Posteriori2').processData(0).cov[0:2,0:2])
            # sat2covTrs[loopIndex+1]= np.trace(gins('Posteriorj2').processData(0).cov[0:2,0:2])
            # timediff[loopIndex+1]  = gins('TimeDiff').processData(0).cov

        # x =  np.linspace(0,loopTime,loopTime+1)
        # plt.figure()
        # plt.plot(x,sat1covTrs,'r',x,sat2covTrs,'g',x,timediff)
        # plt.show()
        # TODO some results output.
        
        results = {}
        results['sat1KineState']  = createStateWrapperFromMsg(gins('Posteriori2').getBeliefs())
        results['sat2KineState']  = createStateWrapperFromMsg(gins('Posteriorj2').getBeliefs())
        results['timeDriftDiff']  = createStateWrapperFromMsg(gins('TimeDiff').getBeliefs()) 
        return results
    
    
    def forwardPass(self):
        logging.debug('I am forward Pass')
        pg  =  self.graph.passMsg
        gins = self.graph.getNodeInstance
        gins('PriorMobilei2').model.setForward()
        gins('PriorMobilej2').model.setForward()

        #Priori
        pg('PriorMobilei1','Posteriori1')
        pg('PriorMobilej1','Posteriorj1')
        pg('TimeDiffPrior','TimeDiff')

        #motional and time diff estimation
        pg('Posteriori1','PriorMobilei2')
        pg('PriorMobilei2','Posteriori2')
        
        pg('Posteriorj1','PriorMobilej2')
        pg('PriorMobilej2','Posteriorj2')

        pg('Posteriori1','MutualProp1')
        pg('Posteriorj1','MutualProp1')
        pg('Posteriori2','MutualProp2')
        pg('Posteriorj2','MutualProp2')

        pg('MutualProp1','TimeDiff')
        pg('MutualProp2','TimeDiff')

        #mutual
        # pg('Posteriori1','MutualProp1') # TODO: seems like redundant
        pg('TimeDiff','MutualProp1')
        pg('TimeDiff','MutualProp2')

        pg('MutualProp1','Posteriorj1')


        pg('Posteriorj1','PriorMobilej2')
        pg('PriorMobilej2','Posteriorj2')
        pg('Posteriorj2','MutualProp2')


        pg('MutualProp2','Posteriori2')

        logging.info('Time diff is mean/std %f/%f '%(gins('TimeDiff').getBeliefs().mean,np.sqrt(gins('TimeDiff').getBeliefs().cov)))

    def backwardPass(self,mode = 1):
        logging.debug('I am backward Pass')
        pg  = self.graph.passMsg
        # backward satellite propagation 

        #stage 2
        pg('Posteriorj1','MutualProp1')
        pg('Posteriori2','MutualProp2')
        pg('MutualProp1','TimeDiff')
        pg('MutualProp2','TimeDiff')

        if mode == 1:
            #stage 3
            gins = self.graph.getNodeInstance
            gins('PriorMobilei2').model.setForward(False)
            gins('PriorMobilej2').model.setForward(False)

            pg('TimeDiff','MutualProp1')
            pg('MutualProp1','Posteriori1')
            pg('TimeDiff','MutualProp2')
            pg('MutualProp2','Posteriorj2')

            pg('Posteriori2','PriorMobilei2')
            pg('PriorMobilei2','Posteriori1')
            pg('Posteriorj2','PriorMobilej2')
            pg('PriorMobilej2','Posteriorj1')
    
if __name__ == '__main__':
    print 'I am running'
    # codes for proving convergence of this factor graph
    